{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_core.messages import HumanMessage, SystemMessage\n",
    "from langgraph.graph import START, END, StateGraph, MessagesState\n",
    "from langgraph.prebuilt import tools_condition, ToolNode\n",
    "from IPython.display import Image, display\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "\n",
    "\n",
    "\n",
    "def multiply_values(a, b):\n",
    "    \"\"\"\n",
    "    Multiply two values and return the result.\n",
    "\n",
    "    Parameters:\n",
    "        a (float): The first value.\n",
    "        b (float): The second value.\n",
    "\n",
    "    Returns:\n",
    "        float: The product of a and b.\n",
    "    \"\"\"\n",
    "    return a * b\n",
    "\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "toolbox = [multiply_values]\n",
    "llm = llm.bind_tools(toolbox)\n",
    "\n",
    "\n",
    "# System message\n",
    "chatbot_system_message = SystemMessage(content=(\"\"\"\n",
    "You are a helpful and knowledgeable chatbot assistant. \n",
    "Your goal is to provide clear and accurate answers to user questions based on the information they provide. \n",
    "Stay focused, concise, and ensure your responses are relevant to the context of the conversation. \n",
    "If you don’t have enough information, ask for clarification.”\n",
    "\"\"\"))\n",
    "\n",
    "# Nodes\n",
    "def chatbot(state: MessagesState):\n",
    "   return {\"messages\": [llm.invoke([chatbot_system_message] + state[\"messages\"])]}\n",
    "\n",
    "def summarization(state: MessagesState):\n",
    "    summary_message = HumanMessage(content=\"\"\"\n",
    "    Summarize the above conversation while preserving full context, key points, and user intent.\n",
    "    Only return the summarized content. Do not add explanations, section headers, or extra commentary.\n",
    "    \"\"\")\n",
    "    \n",
    "    return {\"messages\": [llm.invoke(state[\"messages\"] + [summary_message])]}\n",
    "\n",
    "\n",
    "#Edges\n",
    "def should_summarize(state: MessagesState) -> bool:\n",
    "    return len(state[\"messages\"]) > 2\n",
    "\n",
    "def extended_condition(state: MessagesState):\n",
    "    # First, check if the chatbot indicates a tool call\n",
    "    if tools_condition(state) == \"tools\":\n",
    "        return \"tools\"\n",
    "    \n",
    "    # Then, if no tool is needed, decide if summarization is needed\n",
    "    if should_summarize(state):\n",
    "        return \"summarization\"\n",
    "    \n",
    "    # Otherwise, finish the graph\n",
    "    return END\n",
    "\n",
    "# Graph\n",
    "builder = StateGraph(MessagesState)\n",
    "\n",
    "# Define nodes: these do the work\n",
    "builder.add_node(\"chatbot\", chatbot)\n",
    "builder.add_node(\"tools\", ToolNode(toolbox))\n",
    "builder.add_node(\"summarization\", summarization)\n",
    "\n",
    "builder.add_edge(START, \"chatbot\")\n",
    "builder.add_conditional_edges(\"chatbot\", extended_condition)\n",
    "builder.add_edge(\"tools\", \"chatbot\")\n",
    "builder.add_edge(\"summarization\", END)\n",
    "\n",
    "memory = MemorySaver()\n",
    "graph = builder.compile(checkpointer=memory)\n",
    "\n",
    "# Show\n",
    "display(Image(graph.get_graph(xray=True).draw_mermaid_png()))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {\"configurable\": {\"thread_id\": 1}}\n",
    "messages = [HumanMessage(content=\"What is 2 multiplied by 3?\")]\n",
    "messages = graph.invoke({\"messages\": messages}, config)\n",
    "for m in messages['messages']:\n",
    "    m.pretty_print()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
